package com.dsbj.it.DingTalkBotStarter.config;

import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.stereotype.Component;

import java.sql.Connection;
import java.util.Arrays;
import java.util.Collection;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * MyBatis-Plus SQL注入拦截器
 * 用于拦截和过滤SQL注入攻击
 */
@Slf4j
@Component
@Intercepts({
    @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
    @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
public class SqlInjectionInterceptor implements InnerInterceptor {

    // SQL 注入正则表达式检测
    private static final Pattern SQL_INJECTION_PATTERN = Pattern.compile(
            "(?:')|(?:--)|(/\\*(?:.|[\\n\\r])*?\\*/)|(\\b(select|update|and|or|delete|insert|trancate|char|into|substr|ascii|declare|exec|count|master|into|drop|execute)\\b)",
            Pattern.CASE_INSENSITIVE);

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        // 检查参数中的SQL注入
        BoundSql boundSql = sh.getBoundSql();
        Object parameterObject = boundSql.getParameterObject();

        if (parameterObject != null) {
            // 检查参数值中的SQL注入
            checkParameterForSqlInjection(parameterObject);
        }
    }

    private void checkParameterForSqlInjection(Object parameterObject) {
        if (parameterObject instanceof Map) {
            // 处理Map参数
            ((Map<?, ?>) parameterObject).values().forEach(this::checkSingleValue);
        } else if (parameterObject instanceof Collection) {
            // 处理集合参数
            ((Collection<?>) parameterObject).forEach(this::checkSingleValue);
        } else if (parameterObject.getClass().isArray()) {
            // 处理数组参数
            Arrays.stream((Object[]) parameterObject).forEach(this::checkSingleValue);
        } else {
            // 处理普通对象参数
            checkSingleValue(parameterObject);
        }
    }

    private void checkSingleValue(Object value) {
        if (value != null) {
            String strValue = value.toString();
            if (hasSqlInjection(strValue)) {
                throw new RuntimeException("参数中检测到可能的SQL注入攻击: " + strValue);
            }
        }
    }

    /**
     * 检查是否有SQL注入
     */
    private boolean hasSqlInjection(String value) {
        if (StringUtils.isBlank(value)) {
            return false;
        }
        Matcher matcher = SQL_INJECTION_PATTERN.matcher(value);
        return matcher.find();
    }
}
